from Search_Databases import build_map
from itertools import combinations,product




def get_hits(seq, word_map, word_size):
    '''
    allows at most 1 character to be different
    '''
    def one_edit(word, alphabet = "ATCG"):
        res = {word}
        for i in range(len(word)):
            for editor in alphabet:
                if editor == word[i]:
                    continue
                res.add(word[:i]+editor+word[i+1:])
        return res
    hits = []
    for i in range(len(seq)-word_size+1):
        subseq = seq[i:i+word_size]
        for one_edited in one_edit(subseq):
            if one_edited in word_map:
                tag_indexs = word_map[one_edited]
                hits.extend([(index,i) for index in tag_indexs])
    return hits

def check_mismatch(seq1,seq2,mismatch):
    count = 0
    for c1,c2 in zip(seq1,seq2):
        if c1 != c2:
            count += 1
        if count > mismatch:
            return False
    else:
        return True

def get_approximate(seq,max_mismatch,alphabet = "ATCG"):
    result = []
    alphabet = set(alphabet)
    for mismatch in range(1,max_mismatch+1):
        for charge_pos in combinations(range(len(seq)),mismatch):
            charge_able = [ alphabet.difference(seq[pos]) 
                           for pos in charge_pos]
            charged = seq
            for chars in product(*charge_able):
                for index,char in zip(charge_pos,chars):
                    charged = charged[:index] + char + charged[index+1:]
                result.append(charged)
    return result

def search_uniprot():
    from Bio.Blast import NCBIXML 
    from Bio.Blast import NCBIWWW 
    from Bio import SeqIO

    file = 'bioinformatic_algorithm_colg/5.Search_Databases/O14727.fasta'
    seqrecord = SeqIO.read(open(file), format="fasta") 
    print (len(seqrecord.seq))
    
    result_handle = NCBIWWW.qblast("blastp", "swissprot", seqrecord.format("fasta"))

    record = NCBIXML.read(result_handle)

    print ("PARAMETERS:")
    print ("Database: " + record.database)
    print ("Matrix: " + record.matrix)
    print ("Gap penalties: ", record.gap_penalties)

    nhits = len(record.alignments) 
    print ("number hits: ", nhits)

    res = []
    for alignment in record.alignments:
        evalue = alignment.hsps[0].expect
        accession = alignment.accession
        leng = alignment.hsps[0].align_length
        res.append(accession + " - " + str(evalue) + " length:" + str(leng) )

    print("E-values and length of alignments:")
    for s in res: print(s)

    result_handle2 = NCBIWWW.qblast("blastp", "swissprot",
                                    seqrecord.format("fasta"),
                                    entrez_query = "Saccharomyces cerevisiae[organism]" )

    record2 = NCBIXML.read(result_handle2)

    first_alignment = record2.alignments[0]

    print ("Accession: " + first_alignment.accession)
    print ("Hit id: " + first_alignment.hit_id)
    print ("Definition: " + first_alignment.hit_def)

    print ("Number of HSPs: ", len(first_alignment.hsps))

    for hsp in first_alignment.hsps:
        print ("E-value: ", hsp.expect)
        print ("Length: ", hsp.align_length)
        print ("Identities: ", hsp.identities)
        print ("Query start: ", hsp.query_start)
        print ("Sbjct start: ", hsp.sbjct_start)
        print (hsp.query[0:90])
        print (hsp.match[0:90])
        print (hsp.sbjct[0:90])
        print ("")

if __name__ == "__main__":

    print('### ex1 ###')
    query = 'ATCTCTATCGCAGC'
    seq = 'GGGCGGCCGGCATCTCTATCGCAGCATTATATTCTC'
    word_size = 7
    word_map = build_map(query, word_size)
    print(get_hits(seq, word_map, word_size))

    print('### ex2 ###')
    print(check_mismatch('ATCTATC','ATGTATC',3))
    print(check_mismatch('ATCTATC','ATGTGTC',1))
    print(get_approximate('AAAA',2))
    print(len(get_approximate('AAAA',2)))

    print('### ex3 ###')
    search_uniprot()
    pass